{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\" \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import tensorflow as tf\n",
    "gpu_options = tf.GPUOptions(allow_growth=True)\n",
    "sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))\n",
    "\n",
    "import numpy as np\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "import os.path as osp\n",
    "import random\n",
    "from PIL import Image\n",
    "from scipy import ndimage\n",
    "import tensorflow as tf\n",
    "from spatial_transformer import transformer\n",
    "import numpy as np\n",
    "from tf_utils import weight_variable, bias_variable, dense_to_one_hot\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_STL(path, num_batch):\n",
    "    h = 384\n",
    "    w = 384\n",
    "    im = cv2.imread(path[0])\n",
    "    im = im / 255.\n",
    "    im = cv2.resize(im, (w, h), interpolation=cv2.INTER_CUBIC)\n",
    "    \n",
    "    im = im.reshape(1, h, w, 3)\n",
    "    im = im.astype('float32')\n",
    "    \n",
    "    batch = np.append(im, im, axis=0)\n",
    "    for p in path: \n",
    "        im = cv2.imread(p)\n",
    "        im = im / 255.\n",
    "        im = cv2.resize(im, (w, h), interpolation=cv2.INTER_CUBIC)\n",
    "        im = im.reshape(1, h, w, 3)\n",
    "        im = im.astype('float32')\n",
    "        batch = np.append(batch, im, axis=0)\n",
    "    \n",
    "    batch = batch[2:,:,:,:]\n",
    "\n",
    "    out_size = (h, w)\n",
    "\n",
    "    # %% Simulate batch\n",
    "    x = tf.placeholder(tf.float32, [None, h, w, 3])\n",
    "    x = tf.cast(batch, 'float32')\n",
    "\n",
    "    # %% Create localisation network and convolutional layer\n",
    "    with tf.variable_scope('spatial_transformer_0'):\n",
    "\n",
    "        # %% Create a fully-connected layer with 6 output nodes\n",
    "        n_fc = 6\n",
    "        W_fc1 = tf.Variable(tf.zeros([h * w * 3, n_fc]), name='W_fc1')\n",
    "\n",
    "        # %% Zoom into the image\n",
    "        a, b, c, d, e, f = np.random.random(6)/10\n",
    "\n",
    "        initial = np.array([[1-a, b, c], [d, 1-e, f]])\n",
    "        initial = initial.astype('float32')\n",
    "        initial = initial.flatten()\n",
    "\n",
    "        b_fc1 = tf.Variable(initial_value=initial, name='b_fc1')\n",
    "        h_fc1 = tf.matmul(tf.zeros([num_batch, h * w * 3]), W_fc1) + b_fc1\n",
    "        h_trans = transformer(x, h_fc1, out_size)\n",
    "\n",
    "    # %% Run session\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        y = sess.run(h_trans, feed_dict={x: batch})\n",
    "        sess.close()\n",
    "    \n",
    "    return y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_data_path = \"source data path\"#\"/data4/wangpengxiao/danbooru2017/original\"\n",
    "STL_path = \"STL result path\"#\"/data4/wangpengxiao/danbooru2017/original_STL\"\n",
    "\n",
    "source_img_path = glob.glob(osp.join(source_data_path,'*/*.jpg'))\n",
    "source_img_path += glob.glob(osp.join(source_data_path,'*/*.png'))\n",
    "source_img_path = sorted(source_img_path)\n",
    "\n",
    "batch_size = 16\n",
    "\n",
    "os.makedirs(STL_path,exist_ok=True)\n",
    "q = []\n",
    "count = 0\n",
    "c = 0\n",
    "for path in tqdm(source_img_path):\n",
    "    c += 1\n",
    "    if c != 0 :\n",
    "        if count == batch_size-1 :\n",
    "            q.append(path)\n",
    "            tf.reset_default_graph()\n",
    "            im = get_STL(q, batch_size)\n",
    "            tf.get_default_graph().finalize()\n",
    "            for j in range(len(im)):\n",
    "                img = im[j]\n",
    "                amin, amax = img.min(), img.max() # 求最大最小值\n",
    "                img = (img-amin)/(amax-amin) # (矩阵元素-最小值)/(最大值-最小值)\n",
    "                \n",
    "                cv2.imwrite(osp.join(STL_path, osp.basename(q[j])), (img*255).astype('uint8'))            \n",
    "                \n",
    "            count = 0\n",
    "            q = []\n",
    "        else:\n",
    "            count += 1\n",
    "            q.append(path)\n",
    "    else:\n",
    "        continue\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
